import torch
import pytorch_lightning as pl
from itertools import chain

class FasterRCNN_Lightning(pl.LightningDataModule):
    def __init__(self,
                 model: torch.nn.Module,
                 lr: float=1e-4,
                 iou_threshold: float=0.5
                 ) -> None:
        super().__init__()

        # Model

        self.model = model

        # Classes
        self.num_classes = self.model.num_classes

        # Learning Rate
        self.lr = lr

        # IoU threshold
        self.iou_threshold = iou_threshold

        # Transformation parameters
        self.mean = model.image_mean
        self.std = model.image_std
        self.min_size = model.min_size
        self.max_size = model.max_size

        # Save hyperparameters
        self.save_hyperparameters()

    
    def forward(self, x):
        self.model.eval()
        return self.model(x)
    
    def training_setp(self, batch, batch_idx):
        images, targets = batch
        loss_dict = self.model(images, targets)
        loss = sum(loss for loss in loss_dict.values())

        self.log_dict(loss_dict)

        return loss

    # def validation_step(self, batch, batch_idx):
    #     # Batch
    #     images, targets = batch
    # 
    #     preds = self.model(images)

    def test_step(self, batch, batch_idx):
        images, targets = batch

        preds = self.model(x)

        gt_boxes = []

